import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt

im = Image.open('./Data/cat.png').convert('L')  # 读入一张灰度图的图片
im = np.array(im, dtype='float32')  # 将其转换为一个矩阵
plt.imshow(im.astype('uint8'), cmap='gray')  # 转换为可视化图片并放到画布上
plt.show()  # 展示画布


class simpleCNN(nn.Module):
    def __init__(self):
        super(simpleCNN, self).__init__()
        # 刚放入时图片数据大小为(batch, 3, H, W)
        # batch 每轮处理的图片数  3 表示通道数  H 每张图片的分辨率高  W 每张图片的分辨率宽

        # 开始添加第一层卷积神经网络
        layer1 = nn.Sequential()
        layer1.add_module('conv1', nn.Conv2d(3, 32, 3, 1, padding=1))
        # 第一次卷积后(batch, 32, H, W)
        # Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        # in_channels 输入数据体的深度  out_channels 输出数据体的深度  kernel_size 滤波器(卷积核)的大小  stride 滑动步长  padding 填几圈0
        layer1.add_module('relu1', nn.ReLU(True))
        # 第一次激活后(batch, 32, H, W),激活函数不会改变数据的大小
        layer1.add_module('pool1', nn.MaxPool2d(2, 2))
        # 第一次池化后(batch, 32, H/2, W/2)
        self.layer1 = layer1

        # 开始添加第二层卷积神经网络
        layer2 = nn.Sequential()
        layer2.add_module('conv2', nn.Conv2d(32, 64, 3, 1, padding=1))
        # 第二次卷积后(batch, 64, H/2, W/2)
        layer2.add_module('relu2', nn.ReLU(True))
        layer2.add_module('pool2', nn.MaxPool2d(2, 2))
        # 第二次池化后(batch, 64, H/2/2, W/2/2)
        self.layer2 = layer2

        # 开始添加第三层卷积神经网络
        layer3 = nn.Sequential()
        layer3.add_module('conv3', nn.Conv2d(64, 128, 3, 1, padding=1))
        layer3.add_module('relu3', nn.ReLU(True))
        layer3.add_module('pool3', nn.MaxPool2d(2, 2))
        self.layer3 = layer3

        # 开始添加第四层卷积神经网络
        layer4 = nn.Sequential()
        layer4.add_module('fc1', nn.Linear(2048, 512))
        layer4.add_module('fc_relu1', nn.ReLU(True))
        layer4.add_module('pool4', nn.MaxPool2d(2, 2))
        self.layer4 = layer4
